#include <benchmark/benchmark.h>
#include <algorithm>
#include <array>
#include <cstddef>
#include <immintrin.h>
#include <random>
#include <vector>

// 保证 32 字节对齐，便于 AVX2 访问
template <typename T>
using aligned_vector = std::vector<T, std::allocator<T>>;

constexpr std::size_t kSize = 1 << 20;  // 1,048,576

aligned_vector<float> make_data() {
    aligned_vector<float> data(kSize);
    std::mt19937 rng{123};
    std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
    std::generate(data.begin(), data.end(), [&] { return dist(rng); });
    return data;
}

// 朴素版本
void add_scalar(const float* a, const float* b, float* x, std::size_t size) {
    for (std::size_t i = 0; i < size; ++i) {
        x[i] = a[i] + b[i];
    }
}

// 手工 4 次展开
void add_unrolled4(const float* a, const float* b, float* x, std::size_t size) {
    std::size_t i = 0;
    for (; i + 4 <= size; i += 4) {
        x[i + 0] = a[i + 0] + b[i + 0];
        x[i + 1] = a[i + 1] + b[i + 1];
        x[i + 2] = a[i + 2] + b[i + 2];
        x[i + 3] = a[i + 3] + b[i + 3];
    }
    for (; i < size; ++i) {
        x[i] = a[i] + b[i];
    }
}

// 显式 AVX2 实现（一次处理 8 个 float）
void add_avx2(const float* a, const float* b, float* x, std::size_t size) {
    std::size_t i = 0;
    for (; i + 8 <= size; i += 8) {
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        __m256 vx = _mm256_add_ps(va, vb);
        _mm256_storeu_ps(x + i, vx);
    }
    for (; i < size; ++i) {
        x[i] = a[i] + b[i];
    }
}

// -----------------------------------------------------------------------------
// Google Benchmark 入口
// -----------------------------------------------------------------------------

static void BenchmarkScalar(benchmark::State& state) {
    const auto a = make_data();
    const auto b = make_data();
    auto x       = aligned_vector<float>(kSize);

    for (auto _ : state) {
        add_scalar(a.data(), b.data(), x.data(), kSize);
        benchmark::ClobberMemory();
    }

    const double bytes = static_cast<double>(kSize) * sizeof(float) * 3;
    state.counters["bytes"] = benchmark::Counter(bytes,
                                                 benchmark::Counter::kIsIterationInvariantRate);
}

static void BenchmarkUnrolled4(benchmark::State& state) {
    const auto a = make_data();
    const auto b = make_data();
    auto x       = aligned_vector<float>(kSize);

    for (auto _ : state) {
        add_unrolled4(a.data(), b.data(), x.data(), kSize);
        benchmark::ClobberMemory();
    }

    const double bytes = static_cast<double>(kSize) * sizeof(float) * 3;
    state.counters["bytes"] = benchmark::Counter(bytes,
                                                 benchmark::Counter::kIsIterationInvariantRate);
}

static void BenchmarkAVX2(benchmark::State& state) {
    const auto a = make_data();
    const auto b = make_data();
    auto x       = aligned_vector<float>(kSize);

    for (auto _ : state) {
        add_avx2(a.data(), b.data(), x.data(), kSize);
        benchmark::ClobberMemory();
    }

    const double bytes = static_cast<double>(kSize) * sizeof(float) * 3;
    state.counters["bytes"] = benchmark::Counter(bytes,
                                                 benchmark::Counter::kIsIterationInvariantRate);
}

BENCHMARK(BenchmarkScalar);
BENCHMARK(BenchmarkUnrolled4);
BENCHMARK(BenchmarkAVX2);

BENCHMARK_MAIN();